from settings import *

import shapely
from tqdm import tqdm
from osgeo import gdal

class DEM_insert_tool():
    def __init__(self, table_name, dem_filepath=r"C:\Users\PC\Desktop\路网匹配\会宁数据\栅格数据\会宁县_alosdem_utm.tif"):
        self.table_name = table_name
        self.dem_filepath = dem_filepath
        self.import_dem(self.dem_filepath)  # 导入dem数据

    def import_dem(self, dem_filepath):
        # 导入DEM
        dataset = gdal.Open(dem_filepath)  # 使用gdal打开文件
        # 获取DEM的所有信息
        self.im_width = dataset.RasterXSize  # 栅格矩阵的列数
        self.im_height = dataset.RasterYSize  # 栅格矩阵的行数
        self.im_bands = dataset.RasterCount  # 波段数
        self.im_geotrans = dataset.GetGeoTransform()
        self.im_projection = dataset.GetProjection()
        # 获取DEM数据
        self.im_data = dataset.ReadAsArray(0, 0, self.im_width, self.im_height)
        plt.imshow(self.im_data)

    def insert_dem_Workflow(self, table_name):
        """
        Args:
            table_name: 需要进行dem插值的表
                        必要字段(gid int, qdzh float, zdzh float, geom linestring)
                        新增字段(length_2d float 原始长度, length_zh 桩号长度, length_3d 插值后长度, method_3d text 插值方法)

        Returns:
            table_name

        """
        cur = conn.cursor()
        # 增加信息存储列，并计算2D长度
        sql = f"""
                alter table {table_name} add column if not exists length_2d float ;
                alter table {table_name} add column if not exists length_3d float ;
                alter table {table_name} add column if not exists length_zh float ;
                alter table {table_name} add column if not exists method_3d text ;
                update {table_name} set length_2d = st_length(geom);
                update {table_name} set length_zh = (zdzh - qdzh)*1000;"""
        cur.execute(sql)
        conn.commit()

        gdf = gpd.read_postgis(f"select * from {table_name}", conn, geom_col="geom")

        index_result = list(gdf.columns).index('length_3d')
        index_method = list(gdf.columns).index('method_3d')
        index_zh_len = list(gdf.columns).index('length_zh')
        index_2d_len = list(gdf.columns).index('length_2d')
        index_geometry = list(gdf.columns).index("geom_3d")

        for i in tqdm(range(len(gdf))):
            fea = gdf.iloc[i]
            mile = gdf.iloc[i, index_zh_len]
            Length2D = gdf.iloc[i, index_2d_len]
            ansResult = -1
            resultGeo = fea.geom
            if Length2D < mile:
                # 进行DEM插值
                geo3D_res = self.InsertDEM(fea.geom, self.im_geotrans)
                geo3D = geo3D_res[0]
                using = geo3D_res[1]  # 是否在DEM范围内
                if (using == 1):
                    Length3D = self.getLength3D(geo3D)
                    if Length3D < mile:  # 是否需要看看精度？达不达标？
                        # 进行加密插值
                        ansGeo = self.InsertPois(fea.geom, 20)
                        ansGeo3D = self.InsertDEM(ansGeo, self.im_geotrans, method='BiLinary', noData=-9999)[0]
                        ansResult = self.getLength3D(ansGeo3D)
                        gdf.iloc[i, index_method] = "DEM_Insert"
                        resultGeo = ansGeo3D
                    else:
                        # 保存结果
                        ansResult = Length3D
                        gdf.iloc[i, index_method] = "DEM"
                        resultGeo = geo3D
                else:
                    ansResult = Length2D
                    gdf.iloc[i, index_method] = "outsideDEM2D"
            else:
                # 保存结果
                ansResult = Length2D
                gdf.iloc[i, index_method] = "2D"
            gdf.iloc[i, index_geometry] = resultGeo
            gdf.iloc[i, index_result] = ansResult

        # 插值结果更新到融合结果当中
        conn.commit()
        gdf.to_postgis(table_name, engine, if_exists="replace")

    # 图像与WGS84的互换
    def to_WGS84(self, x, y, geos):
        geox = geos[0] + geos[1] * x + geos[2] * y
        geoy = geos[3] + geos[4] * x + geos[5] * y
        return [geox, geoy]

    def to_xy(self, p, geos):
        lng = p[0]
        lat = p[1]
        x = np.array([[geos[1], geos[2]], [geos[4], geos[5]]])
        x1 = np.linalg.inv(x)
        y = np.array([[lng - geos[0]], [lat - geos[3]]])
        return np.dot(x1, y)

    def Nearest(self, data, x, y):
        x_int = int(x + 0.5)
        y_int = int(y + 0.5)
        return data[y_int][x_int]

    def BiLinary(self, data, x, y):
        x_int = int(x)
        y_int = int(y)
        # data的访问是先列后行
        return data[y_int, x_int] * (1 - (x - x_int)) * (1 - (y - y_int)) + data[y_int, x_int + 1] * (
                    1 - (x_int + 1 - x)) * (1 - (y - y_int)) + data[y_int + 1, x_int] * (1 - (x - x_int)) * (
                           1 - (y_int + 1 - y)) + data[y_int + 1, x_int + 1] * (1 - (x_int + 1 - x)) * (
                           1 - (y_int + 1 - y))

    # 两点距离
    def Distance(self, p1, p2):
        return np.sqrt(sum((np.array(p1) - np.array(p2)) ** 2))

    def InsertDEM(self, geo, im_geotrans, method='BiLinary', noData=-1):
        """
        向线段中内插DEM的高程数据
        """
        using = 1
        res_geo = geo
        try:
            g_list = list(geo.coords)
            temp_list = []
            tag = True
            for p in g_list:
                dem_xy = self.to_xy(p, im_geotrans)
                if (dem_xy[0] > self.im_width - 1 or dem_xy[1] > self.im_height - 1 or dem_xy[0] < 0 or dem_xy[1] < 0):
                    using = 0
                    break
                if (method == 'BiLinary'):
                    dem_height = self.BiLinary(self.im_data, *dem_xy)[0]
                elif (method == 'Nearest'):
                    dem_height = self.Nearest(self.im_data, *dem_xy)
                if (dem_height == self.noData or dem_height < 0):
                    # print("dem_height:{0}".format(dem_height))
                    dem_height = -1
                    using = 0
                    tag = False
                temp_list.append([p[0], p[1], dem_height])
            res_geo = shapely.geometry.LineString(temp_list)
            if (not tag):
                using = 0
        except NotImplementedError as e:
            using = 0
            print("multiPart")
        return [res_geo, using]

    def InsertPois(self, geo, max_dis):
        """
        内插点，max_dis为间隔
        """
        tag = True
        pList = list(geo.coords)
        while (tag):
            tag = False
            for j in range(len(pList) - 1):
                s = pList[j]
                e = pList[j + 1]
                dis = self.Distance(s, e)
                if dis > max_dis:
                    # 考虑在两点之间内插
                    pList.insert(j + 1, (np.array(s) + np.array(e)) / 2)
                    tag = True
                    break
        return shapely.geometry.LineString(pList)

    def getLength3D(self, geo):
        su = 0
        pList = list(geo.coords)
        for i in range(len(pList) - 1):
            s = pList[i]
            e = pList[i + 1]
            su += self.Distance(s, e)
        return su

